package Leetcode.图;

import java.util.Arrays;

public class KruskalAlgorithm {
    static class Edge implements Comparable<Edge> {
        int src, dest, weight;

        public int compareTo(Edge compareEdge) {
            return this.weight - compareEdge.weight;
        }
    }

    static class Subset {
        int parent, rank;
    }

    static int find(Subset[] subsets, int i) {
        if (subsets[i].parent != i)
            subsets[i].parent = find(subsets, subsets[i].parent);
        return subsets[i].parent;
    }

    static void union(Subset[] subsets, int x, int y) {
        int xroot = find(subsets, x);
        int yroot = find(subsets, y);

        if (subsets[xroot].rank < subsets[yroot].rank)
            subsets[xroot].parent = yroot;
        else if (subsets[xroot].rank > subsets[yroot].rank)
            subsets[yroot].parent = xroot;
        else {
            subsets[yroot].parent = xroot;
            subsets[xroot].rank++;
        }
    }

    public static void kruskal(int[][] graph, int V) {
        Edge[] result = new Edge[V]; // MST will have V-1 edges
        int e = 0; // An index variable, used for result[]
        int[][] edges = new int[V * V][3]; // Array to store edges in the graph
        int edgeCount = 0;

        // 初始化图
        for (int i = 0; i < V; i++) {
            for (int j = i + 1; j < V; j++) {
                if (graph[i][j] != 0 && graph[i][j] != Integer.MAX_VALUE) {
                    edges[edgeCount][0] = i;
                    edges[edgeCount][1] = j;
                    edges[edgeCount][2] = graph[i][j];
                    edgeCount++;
                }
            }
        }

        // 边的权值数组，并进行升序排序
        Edge[] edgeArray = new Edge[edgeCount];
        for (int i = 0; i < edgeCount; i++) {
            edgeArray[i] = new Edge();
            edgeArray[i].src = edges[i][0];
            edgeArray[i].dest = edges[i][1];
            edgeArray[i].weight = edges[i][2];
        }
        Arrays.sort(edgeArray);

        Subset[] subsets = new Subset[V];
        for (int i = 0; i < V; i++) {
            subsets[i] = new Subset();
            subsets[i].parent = i;
            subsets[i].rank = 0;
        }

        for (Edge edge : edgeArray) {
            // 按权重顺序选择下一个边
            int x = find(subsets, edge.src);
            int y = find(subsets, edge.dest);

            // 如果包括这条边不会形成环，则将其包括在结果中
            if (x != y) {
                result[e++] = edge;
                union(subsets, x, y);
                if (e == V - 1) break; // MST has V-1 edges
            }
        }

        System.out.println("Following are the edges in the constructed MST");
        for (int i = 0; i < e; i++) {
            System.out.println(result[i].src + " -- " + result[i].dest + " == " + result[i].weight);
        }
    }

    public static void main(String[] args) {
        int V = 4; // Number of vertices in the graph
        int[][] graph = new int[][]{
            {0, 10, 6, 5},
            {10, 0, 0, 15},
            {6, 0, 0, 4},
            {5, 15, 4, 0}
        };

        kruskal(graph, V);
    }
}
